#include "include/AdvectionDiffusion.h"
#include <fstream>

using namespace std;

#define PI 3.141592653589793

#define Dim 2

int times = 3;

// 记录误差
Vector err(times);
int ind = 0;

Real Cr = 1;
Real aV = 0.1;

// 对流常数
Real nu = 0.001;

// 速度场
Vector u(Vector v)
{
    Real x = v[0];
    Real y = v[1];
    return {aV * sin(PI * x) * sin(PI * x) * sin(2 * PI * y),
            -aV * sin(2 * PI * x) * sin(PI * y) * sin(PI * y)};
}

Real K[2] = {2 * PI, 4 * PI};

// phi(x,t)
Vector phi(Vector x, Real t)
{
    return Vector({sin(K[0] * x[0] - t) * sin(K[1] * x[1] - t)});
}

// f(x,t)
Vector f(Vector x, Real t)
{
    Real sum1 = 0;
    for (int d = 0; d < Dim; d++)
    {
        sum1 += nu * K[d] * K[d];
    }
    // G<phi>
    sum1 = sum1 * (Real)phi(x, t);

    Real sum2 = 2 * aV * PI *
                (sin(PI * x[0]) * cos(PI * x[0]) * sin(2 * PI * x[1]) -
                 sin(PI * x[1]) * cos(PI * x[1]) * sin(2 * PI * x[0])) *
                (Real)phi(x, t);
    Vector v = u(x);
    Real sum3 = v[0] * K[0] * cos(K[0] * x[0] - t) * sin(K[1] * x[1] - t) +
                v[1] * K[1] * sin(K[0] * x[0] - t) * cos(K[1] * x[1] - t);

    // 计算 phi_t
    Real sum4 = -cos(K[0] * x[0] - t) * sin(K[1] * x[1] - t) -
                sin(K[0] * x[0] - t) * cos(K[1] * x[1] - t);

    return Vector({sum1 + sum2 + sum3 + sum4});
}

Real T0 = 1;

Grid<Dim> test(int N, fstream &fp)
{
    // 空间和时间步长
    Real h = 1.0 / N;
    Real k = Cr * h;

    int SN = N;
    int TN = T0 / k;
    Real T = k * TN;

    // 创建单位区域
    Domain<Dim> domain;

    // 积分算子
    Integrator<3, Dim> I;

    // 包装初值
    FuncX Phi = [=](Vector x)
    {
        return phi(x, 0);
    };

    // 初始化初值网格
    Grid<Dim> U0(SN + 1);
    I.fill(domain, U0, Phi);

    // 填充速度
    vector<Grid<Dim>> V(Dim, Grid<Dim>(SN + 1));
    I.fill(domain, V, u);

    // 初始化填充器
    GhostFiller<Dim> ghostfiller;

    // 创建求解器
    AdvectionDiffusionSolver<Dim> adsolver(f, k, nu, ghostfiller);

    // 开始计算
    Real t = adsolver.solve(U0, V, TN);

    // 终值
    Grid<Dim> U1(SN + 1);
    I.fill(domain, U1, Phi);

    // 理查德外推计算终值
    adsolver.exact(U1, V, TN);

    fp << "SN: " << SN << endl;
    fp << "Time: " << t << endl;

    // 记录网格误差
    err[ind] = TensorOp::grid(U1 - U0);
    fp << "Error: " << err[ind] << endl;
    fp << endl;
    ind++;

    return U0;
}

int main()
{
    int N = 16;

    fstream fp("test2", ios::out);
    fstream fmat("test2.m", ios::out);

    // 记录每次的解
    GridX<Dim> U(times);

    // 测试并记录误差
    for (int i = 0; i < times; i++)
    {
        U[i] = test(N, fp);

        // 输出绘图程序
        Domain<Dim> domain;

        fmat << "XY = [" << endl;

        // 循环输出坐标
        for (int j = 0; j < N + 1; j++)
        {
            Real h = 1.0 / (N + 1);
            fmat << h * j << " " << h * j << endl;
        }

        fmat << "];" << endl;

        fmat << "[X, Y] = meshgrid(XY(:,1), XY(:,2));" << endl;
        fmat << "U = [" << U[i] << "];" << endl;
        fmat << "figure;" << endl;
        fmat << "plot3(X,Y,U,'.b','MarkerSize',5);" << endl;
        fmat << "title('N = " << N << "');" << endl;

        N *= 2;
    }

    // 计算收敛阶
    Real rate = TensorOp::rate(err);

    fp << "Convergence rate: " << rate << endl;

    // 误差向量取对数
    err = TensorOp::log(err);

    fmat << "E = [" << err << "];" << endl;
    fmat << "figure;" << endl;
    fmat << "plot(E,'o-');" << endl;
    fmat << "title('Convergence rate');" << endl;

    fp.close();
    fmat.close();

    return 0;
}